import os
import cv2
import numpy as np
import open3d as o3d
from utils.utils import shift_mask, get_color_map
from utils.read_data import read_pose_from_json


def get_bbox(mask):
    """
    input: 2D mask
    output: object bbox: [min row, max row, min column, max column]
    """

    color_map = get_color_map(N=10)
    c = color_map[1]  # object color
    object_mask = (mask[..., 0] == c[0]) & (mask[..., 1] == c[1]) & (mask[..., 2] == c[2])
    rs, cs = np.where(object_mask)
    return [np.min(rs), np.max(rs), np.min(cs), np.max(cs)]


def get_table_depth(depth, object_bbox, r_width=30, c_width=30, r_attach=5, c_attach=5):
    """
    input: depth, object bbox, ring row&column width
    output: depth (only contains the ring area)
    """
    depth = np.float32(depth)
    r_ids, c_ids = np.where(depth > -1)  # index arrays
    r_ids = r_ids.reshape(depth.shape)
    c_ids = c_ids.reshape(depth.shape)
    valid = ((object_bbox[0] - r_width <= r_ids) & (r_ids <= object_bbox[1] + r_width) & (object_bbox[2] - c_width <= c_ids) & (c_ids <= object_bbox[3] + c_width)) \
        & (~((object_bbox[0] - r_attach <= r_ids) & (r_ids <= object_bbox[1] + r_attach) & (object_bbox[2] - c_attach <= c_ids) & (c_ids <= object_bbox[3] + c_attach)))
    depth[~valid] = 0
    depth = o3d.geometry.Image(depth)
    return depth


def get_distance(points, plane):
    """
    input: points: (N, 3), plane: [A, B, C](Ax+By+Cz=1)
    output: distance from point to plane
    """
    points = np.float32(points)
    plane = np.float32(plane)
    Ds = 1 - points @ plane.reshape(3, 1)
    dists = np.abs(Ds / np.linalg.norm(plane))
    return dists


def single_plane_fitting(points, supervised_normal=None):
    """
    input: point cloud, supervised normal
    output: the table plane parameters
    method: (RANSAC) + Least Square(or PCA), Least Square: Ax+By+Cz=1
    """
    assert points.shape[0] > 100

    if not supervised_normal is None:
        supervised_normal = np.float32(supervised_normal).reshape(3)
    
    # RANSAC
    eps = 0.01
    N = points.shape[0]
    best_plane = {"plane": None, "valid_cnt": -1}
    for k in range(200):  # 200 trials
        ids = np.random.choice(N, 5, replace=False)  # sample 5 points
        pts = points[ids]
        # Least Square
        A = pts.copy()
        b = np.ones((A.shape[0], 1))
        plane = np.linalg.solve(A.T @ A, A.T @ b)  # plane: [A, B, C]
        # check normal
        if not supervised_normal is None:
            similarity = np.sum(plane.reshape(3) * supervised_normal)
            if np.abs(similarity) < 0.9:  # angle > 26 degree
                continue
        # update answer
        dists = get_distance(points, plane)
        valid_cnt = np.sum(dists < eps)
        if valid_cnt > best_plane["valid_cnt"]:
            best_plane = {"plane": plane, "valid_cnt": valid_cnt}
    
    # finalize result
    assert best_plane["valid_cnt"] > 0
    ori_plane = best_plane["plane"]
    dists = get_distance(points, ori_plane)
    valid, _ = np.where(dists < eps)
    A = points[valid].copy()
    print(A.shape)
    A = A[np.random.choice(A.shape[0], 100, replace=True)]
    print(A.shape)
    b = np.ones((A.shape[0], 1))
    final_plane = np.linalg.solve(A.T @ A, A.T @ b)  # plane: [A, B, C]

    # compute inlier number
    dists = get_distance(points, final_plane)
    valid_cnt = np.sum(dists < eps)

    print("N = {}, valid_cnt = {}".format(N, valid_cnt))
    return final_plane


def table_ground_fitting(video_dir, intrinsic_dir):
    """
    input: RGB-D image of frame 0, 2D object mask(or bbox)
    output: table plane and ground plane in the coordinate system of frame 0
    method: sample points near the object -> single plane fitting
    """

    # read data
    p1 = video_dir.find("ZY")
    p2 = video_dir[p1:].find("/")
    ZY = video_dir[p1:p1+p2]
    if ZY == "ZY20210800001":
        rgb_path = os.path.join(video_dir, "align_image", "0.jpg")
    else:
        rgb_path = os.path.join(video_dir, "shift_rgb", "0.jpg")
    depth_path = os.path.join(video_dir, "align_depth", "0.png")
    mask_path = os.path.join(video_dir, "refine_2Dseg", "mask", "00000.png")  # unshifted
    rgb = o3d.io.read_image(rgb_path)
    depth = o3d.io.read_image(depth_path)
    mask = cv2.imread(mask_path, cv2.IMREAD_UNCHANGED)[:, :, ::-1]  # unshifted
    mask = shift_mask(mask, ZY)  # shifted
    intrinsic = np.load(os.path.join(intrinsic_dir, ZY, "intrin.npy"))
    intrinsic_o3d = o3d.camera.PinholeCameraIntrinsic(1920, 1080, intrinsic[0, 0], intrinsic[1, 1], intrinsic[0, 2], intrinsic[1, 2])
    pose_path = os.path.join(video_dir, "objpose", "0.json")
    if not os.path.isfile(pose_path):
        pose_path = os.path.join(video_dir, "objpose", "00000.json")
    object_pose = read_pose_from_json(pose_path)
    
    object_bbox = get_bbox(mask)
    table_depth = get_table_depth(depth, object_bbox, r_width=40, c_width=40, r_attach=10, c_attach=10)

    scene_rgbd = o3d.geometry.RGBDImage.create_from_color_and_depth(rgb, depth, convert_rgb_to_intensity=False)
    scene_pcd = o3d.geometry.PointCloud.create_from_rgbd_image(scene_rgbd, intrinsic_o3d, extrinsic=np.eye(4))
    scene_pcd = scene_pcd.voxel_down_sample(0.01)
    scene_points = np.float32(scene_pcd.points)
    colors = np.float32([0, 0, 1]).reshape(1, 3).repeat(scene_points.shape[0], axis=0)

    # table fitting
    table_rgbd = o3d.geometry.RGBDImage.create_from_color_and_depth(rgb, table_depth, convert_rgb_to_intensity=False)
    table_pcd = o3d.geometry.PointCloud.create_from_rgbd_image(table_rgbd, intrinsic_o3d, extrinsic=np.eye(4))
    o3d.io.write_point_cloud("./table.ply", table_pcd)
    table_points = np.float32(table_pcd.points)
    table_plane = single_plane_fitting(table_points, supervised_normal=object_pose["rotation"][:, 2:])

    # visualization
    eps = 0.01
    dists = get_distance(scene_points, table_plane).reshape(-1)
    colors[dists < eps] = np.float32([1, 0, 0])

    # ground fitting
    camera_ground = 1.8  # kinect_v2 to ground = 170cm
    init_ground_D = camera_ground * np.linalg.norm(table_plane)  # init ground: Ax+By+Cz=init_ground_D
    near = 0.3  # camera_ground += 30cm
    init_ground_plane = table_plane / init_ground_D  # init ground: A'x+B'y+C'z=1
    dists = get_distance(scene_points, init_ground_plane).reshape(-1)
    ground_points = scene_points[dists < near]
    if ground_points.shape[0] >= 20:
        ground_D = np.mean(ground_points @ init_ground_plane.reshape(3, 1))  # ground: A'x+B'y+C'z=ground_D
    else:
        ground_D = 1
    ground_plane = init_ground_plane / ground_D  # ground: A''x+B''y+C''z=1

    # visualization
    eps = 0.01
    dists = get_distance(scene_points, ground_plane).reshape(-1)
    colors[dists < eps] = np.float32([0, 1, 0])
    scene_pcd.colors = o3d.utility.Vector3dVector(colors)
    o3d.io.write_point_cloud("result.ply", scene_pcd)

    return table_plane, ground_plane


if __name__ == "__main__":
    video_dir = "/share/datasets/HOI4D_overall/ZY20210800001/H1/C12/N40/S165/s04/T1"
    # video_dir = "/share/datasets/HOI4D_overall/ZY20210800001/H1/C2/N24/S214/s04/T1"
    # video_dir = "/share/datasets/HOI4D_overall/ZY20210800004/H4/C13/N28/S184/s01/T4"
    
    intrinsic_dir = "/share/datasets/HOI4D_intrinsics"
    table_plane, ground_plane = table_ground_fitting(video_dir, intrinsic_dir)
